# 第四章可变劳动的贝尔曼方程

bellman_ch4 <- function(a0, a1, agrid, vold, thetax, deltax, betax, A){

  # 通过一阶条件获得先求解劳动
  chfun <- function(h){
    c0 <- a0^thetax*h^(1-thetax) - a1 + (1-deltax)*a0
    A*c0/(1-h)-(1-thetax)*a0^thetax*h^(-thetax)
  }
  h0 <- uniroot.all(chfun, interval = c(0,1))
  # print(h0)

  c0 <- a0^thetax*h0^(1-thetax) - a1 + (1-deltax)*a0

  if (c0 < 0) return(-1e10)
  if (a1 > agrid[length(agrid)]) return(a1^2*(-1e10))

  return(log(c0) + A*log(1-h0) + betax * interp1(agrid, vold, xi = a1))
}


# 多个参数压缩成一个
para_cmp <- function(a0, agrid, vold, thetax, deltax, betax, A){
  function(a1){
    bellman_ch4(a0 = a0, a1 = a1, agrid = agrid, vold = vold, thetax = thetax,
                deltax = deltax, betax = betax, A = A)
  }
}


# 包含边界的最大值搜索
findmax <- function(pos, eps, agrid, vold, bellman){
  if (pos$ax == pos$bx){
    pos$bx <- pos$ax + eps * (agrid[2] - agrid[1])
    if (interp1(agrid, vold, pos$bx) < interp1(agrid, vold, pos$ax)){
      aopt <- agrid[1]
    }else{
      # aopt <- GoldenSectionMax(bellman, ay = pos$ax, by = pos$bx,
      #                             cy = pos$cx)
      # browser()
      aopt <- optimize(bellman, c(pos$ax, pos$cx),maximum = T)$maximum
    }
  }else if (pos$bx == pos$cx){
    pos$bx <- pos$cx - eps * (agrid[length(agrid)] - agrid[length(agrid)-1])
    if (interp1(agrid, vold,pos$bx) < interp1(agrid, vold,pos$cx)){
      aopt <- agrid[length(agrid)]
    }else {
      # aopt <- GoldenSectionMax(bellman, ay = pos$ax, by = pos$bx,
      #                             cy = pos$cx)
      aopt <- optimize(bellman, c(pos$ax, pos$cx),maximum = T)$maximum
    }
  }else {
    # aopt <- GoldenSectionMax(bellman, ay = pos$ax, by = pos$bx,
    #                             cy = pos$cx)
    aopt <- optimize(bellman, c(pos$ax, pos$cx),maximum = T)$maximum
  }
  return(aopt)
}
